import os
import cv2
import numpy as np
import pandas as pd
import torch
from torch.utils.data.dataset import Dataset
from torch.autograd import Variable


class ImgDataset(Dataset):
    def __init__(self, data_path, gt_path, gt_downsample, is_training):
        self.data_path = data_path
        self.gt_path = gt_path
        self.gt_downsample = gt_downsample
        self.is_training = is_training
        self.data_files = [filename for filename in os.listdir(data_path)
                           if os.path.isfile(os.path.join(data_path, filename))]

    def __getitem__(self, index):
        img = cv2.imread(os.path.join(self.data_path, self.data_files[index]), 0)
        img = img.astype(np.float32, copy=False)
        # img = cv2.resize(img, (256, 256))
        den = pd.read_csv(os.path.join(self.gt_path, os.path.splitext(self.data_files[index])[0] + '.csv'), sep=',',
                          header=None).as_matrix()
        den = den.astype(np.float32, copy=False)
        ht = img.shape[0]
        wd = img.shape[1]
        ht_1 = 256
        wd_1 = 256
        img = cv2.resize(img, (wd_1, ht_1))
        img = img.reshape((1, 1, img.shape[0], img.shape[1]))
        if self.gt_downsample:
            wd_1 = int(wd_1/8)
            ht_1 = int(ht_1/8)
            den = cv2.resize(den, (wd_1, ht_1))
            den = den * ((wd * ht) / (wd_1 * ht_1))
        else:
            den = cv2.resize(den, (wd_1, ht_1))
            den = den * ((wd * ht) / (wd_1 * ht_1))
        den = den.reshape((1, 1, den.shape[0], den.shape[1]))
        img = img / 255.0
        img = img * 2.0 - 1.0
        den = den * 2.0 - 1.0
        if self.is_training:
            img = Variable(torch.from_numpy(img).type(torch.FloatTensor))
            den = Variable(torch.from_numpy(den).type(torch.FloatTensor))
        else:
            img = Variable(torch.from_numpy(img).type(torch.FloatTensor), requires_grad=False)
            den = Variable(torch.from_numpy(den).type(torch.FloatTensor), requires_grad=False)
        fname = self.data_files[index]
        blob = {'img_data': img, 'gt_data': den, 'file_name': fname}
        return blob

    def __len__(self):
        return len(self.data_files)

    def size(self):
        return self.__len__()


class TestDataset():
    def __init__(self, data_path):
        self.data_path = data_path
        self.data_files = [filename for filename in os.listdir(data_path)
                           if os.path.isfile(os.path.join(data_path, filename))]

    def __getitem__(self, index):
        img = cv2.imread(os.path.join(self.data_path, self.data_files[index]), 0)
        img = img.astype(np.float32, copy=False)
        img = img.reshape((1, 1, img.shape[0], img.shape[1]))
        img = Variable(torch.from_numpy(img).type(torch.FloatTensor))
        fname = self.data_files[index]
        blob = {'img_data': img, 'file_name': fname}
        return blob

    def __len__(self):
        return len(self.data_files)

    def size(self):
        return self.__len__()
